import React, { useState, useEffect } from "react";
import { Tree } from "antd";
import type { DataNode } from "antd/es/tree";

interface Props {
	treeData: DataNode[];
	onChange?: (selectedKeys: React.Key[]) => void;
	initialValues: React.Key[];
}

interface NewDataNode {
	key: string | number;
	title: string | number;
	isSelect?: boolean;
	children?: NewDataNode[];
}

const setDataSelect = (treeData: NewDataNode[], initialValues: React.Key[]) => {
	treeData.forEach(item => {
		if (initialValues.includes(item.key)) {
			item.isSelect = true;
		}
		if (item.children) {
			setDataSelect(item.children, initialValues);
		}
	});
};

const clearNoFillParentId = (treeData: NewDataNode[], initialValues: React.Key[]) => {
	treeData.forEach(item => {
		if (item.children) {
			clearNoFillParentId(item.children, initialValues);
			if (item.isSelect && item.children.filter(item => item.isSelect).length !== item.children.length) {
				let index = initialValues.indexOf(item.key);
				if (index > -1) {
					initialValues.splice(index, 1);
				}
			}
		}
	});
};

const queryChildrenIdIsFill = (treeData: DataNode[], initialValues: React.Key[]) => {
	const data = [...treeData] as NewDataNode[];
	setDataSelect(data, initialValues);
	clearNoFillParentId(data, initialValues);
};

const CustomTreeSelect: React.FC<Props> = ({ treeData, onChange, initialValues }) => {
	const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);
	const [checkedKeys, setCheckedKeys] = useState<React.Key[]>([]);
	const [selectedKeys, setSelectedKeys] = useState<React.Key[]>([]);
	const [autoExpandParent, setAutoExpandParent] = useState<boolean>(true);

	useEffect(() => {
		queryChildrenIdIsFill(treeData, initialValues);
		setCheckedKeys(initialValues);
	}, [initialValues]);

	const onExpand = (expandedKeysValue: React.Key[]) => {
		console.log("onExpand", expandedKeysValue);
		// if not set autoExpandParent to false, if children expanded, parent can not collapse.
		// or, you can remove all expanded children keys.
		setExpandedKeys(expandedKeysValue);
		setAutoExpandParent(false);
	};

	const onCheck = (checkedKeysValue: React.Key[], e: any) => {
		console.log(checkedKeysValue);
		setCheckedKeys(checkedKeysValue);

		setCheckedKeys([...checkedKeysValue]);
		onChange?.([...checkedKeysValue, ...e.halfCheckedKeys]);
	};

	const onSelect = (selectedKeysValue: React.Key[], info: any) => {
		console.log("onSelect", info);
		setSelectedKeys(selectedKeysValue);
	};

	return (
		<Tree
			// @ts-ignore
			onCheck={onCheck}
			checkable
			onExpand={onExpand}
			expandedKeys={expandedKeys}
			autoExpandParent={autoExpandParent}
			defaultCheckedKeys={initialValues}
			checkedKeys={checkedKeys}
			onSelect={onSelect}
			selectedKeys={selectedKeys}
			treeData={treeData}
		/>
	);
};

export default CustomTreeSelect;
